package main

import (
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"

	"github.com/peterh/liner"

	. "gitee.com/shuohe/keytool/msg"

	"gitee.com/shuohe/xutils/xcrypto/base58"
	"gitee.com/shuohe/xutils/xerror"
	"gitee.com/shuohe/xutils/xnet"
	"gitee.com/shuohe/xutils/xterm"
	"gitee.com/shuohe/xutils"
    "gitee.com/shuohe/xutils/xcrypto"
)

var g_chnl 		*xnet.TCPChannel

var errNotEnoughArguments = xerror.SimpleNew("Not enough arguments.")

//=============================================================================
func normalizeKeyAlgo(algo string) string {
	algo = strings.ToUpper(algo)
	switch algo {
	case "AES":
		algo = "AES-128"
	case "ECC":
		algo = "ECC-256"
	case "ED":
		algo = "ED25519"
	}

	return algo
}

func talkToServer(ch *xnet.TCPChannel, req *Bitmap32Message) (*Bitmap32Message, error) {
	err := ch.Write(req.Marshal())
	if err != nil {
		return nil, xerror.SimpleWrap(err, "Failed to send message to the server.")
	}

	rspbytes, err := ch.Read()
	if err != nil {
		return nil, xerror.SimpleWrap(err, "Failed to read message from the server.")
	}

	rsp, err := Bitmap32Message_Unmarshal(rspbytes)
	if err != nil {
		return nil, xerror.SimpleWrap(err, "Can not unmarshal received data.")
	}

	return rsp, nil
}

/**
 * 用法：
 *     genkey <keyAlgo> [--save] [--echo]
 * 说明：
 *     生成指定密码算法的密钥或密钥对。
 * 选项：
 *     --save
 *         将密钥存储在服务端 Keyset 中，返回密钥索引。
 *     --echo
 *         将密钥经过服务端主密钥加密后返回。若没有指定 --save 选项，则本选项默认打开。
 */
func genkey(mode int, args []string) error {
	optSave := false
	optEcho := true

	if len(args) < 2 {
		return errNotEnoughArguments
	}
	if len(args) >= 3 {
		if args[2] == "--save" {
			optSave = true
			optEcho = false
		} else if args[2] != "--echo" {
			return xerror.SimpleNewf("Unknown argument or option '%s'.", args[2])
		}
	}
	if len(args) >= 4 {
		if args[3] == "--save" {
			// args[2] must be '--echo'
			optSave = true
			optEcho = true
		} else if args[3] == "--echo" {
			optEcho = true
		} else {
			return xerror.SimpleNewf("Unknown argument or option '%s'", args[3])
		}
	}

	var kindex int64
	var privkey, pubkey, seckey []byte

	if mode == RunMode_Client {
        options := [2]byte{'0', '1'}
        if optSave {
            options[0] = '1'
        }
        if !optEcho {
            options[1] = '0'
        }

        req := &Bitmap32Message{}
        req.SetString(MsgType, MsgType_GenkeyRequest)
        req.SetString(KeyAlgo, normalizeKeyAlgo(args[1]))
        req.SetString(Options, xutils.Bytes2String(options[:]))
        rsp, err := talkToServer(g_chnl, req)
        if err != nil {
            return err
        }

        // ASSERT: rsp.MsgType == MsgType_GenkeyResponse
        rc := rsp.GetString(RetCode)
        if rc != RetCode_Ok {
            return xerror.SimpleNewf("Server execution failed. Return code: %s.", rc)
        }

        kindex = rsp.GetInt(KeyIndex1)
        privkey = rsp.GetBytes(PrivateKey)
        pubkey = rsp.GetBytes(PublicKey)
        seckey = rsp.GetBytes(SecretKey1)

        xterm.Green().Println("Server execution succeeded.")
    } else { // RunMode_Console
        k, err := xcrypto.GenerateKey(normalizeKeyAlgo(args[1]))
        if err != nil {
            return xerror.SimpleWrapf(err, "Failed to generate key.")
        }

        if optSave {
            g_maxkidx++
            g_keyset.Add(strconv.Itoa(g_maxkidx), k)
            g_keyset.Save(g_keysetPath)
            kindex = int64(g_maxkidx)
        }

        if k.Type() == xcrypto.KeyType_Private {
            pubkey = k.(xcrypto.IPrivateKey).Public().Export()
        }

        if optEcho {
            kbytes := g_mkey.Encrypt(k.Export())
            if k.Type() == xcrypto.KeyType_Private {
                privkey = kbytes
            } else {
                seckey = kbytes
            }
        }
    }

    if kindex > 0 {
        xterm.Blue().Println("Key Index:")
        xterm.Yellow().Println("   ", kindex)
    }
    if privkey != nil && len(privkey) > 0 {
        xterm.Blue().Println("Encrypted Private Key:")
        xterm.Yellow().Println("   ", base58.Encode(privkey))
    }
    if pubkey != nil && len(pubkey) > 0 {
        xterm.Blue().Println("Public Key:")
        xterm.Yellow().Println("   ", base58.Encode(pubkey))
    }
    if seckey != nil && len(seckey) > 0 {
        xterm.Blue().Println("Encrypted Secret Key:")
        xterm.Yellow().Println("   ", base58.Encode(seckey))
    }
    xterm.Reset()

	return nil
}

/**
 * 用法：
 *     getkey <keyIndex>
 * 说明：
 *     返回指定密钥索引对应的密钥值（用主密钥加密后）。
 */
func getkey(mode int, args []string) error {
	if len(args) < 2 {
		return errNotEnoughArguments
	}

	idx, err := strconv.Atoi(args[1])
	if err != nil {
		return xerror.SimpleWrap(err, "Invalid arguments.")
	}

	var privkey, pubkey, seckey []byte

	if mode == RunMode_Client {
        req := &Bitmap32Message{}
        req.SetString(MsgType, MsgType_GetkeyRequest)
        req.SetString(Options, "00")
        req.SetInt(KeyIndex1, int64(idx))
        rsp, err := talkToServer(g_chnl, req)
        if err != nil {
            return err
        }

        // ASSERT: rsp.MsgType == MsgType_GetkeyResponse
        rc := rsp.GetString(RetCode)
        if rc != RetCode_Ok {
            return xerror.SimpleNewf("Server execution failed. Return code: %s.", rc)
        }

        privkey = rsp.GetBytes(PrivateKey)
        pubkey = rsp.GetBytes(PublicKey)
        seckey = rsp.GetBytes(SecretKey1)

        xterm.Green().Println("Server execution succeeded.")
    } else {
        item, err := g_keyset.Get(strconv.Itoa(idx))
        if err != nil {
            return xerror.Wrap(err, "Can not get key from the keyset.")
        }
        if item.Key.Type() == xcrypto.KeyType_Private {
            privkey = g_mkey.Encrypt(item.Key.Export())
            pubkey = item.Key.(xcrypto.IPrivateKey).Public().Export()
        } else {
            seckey = g_mkey.Encrypt(item.Key.Export())
        }
    }

    if privkey != nil && len(privkey) > 0 {
        xterm.Blue().Println("Encrpted Private Key:")
        xterm.Yellow().Println("   ", base58.Encode(privkey))
        xterm.Blue().Println("Public Key:")
        xterm.Yellow().Println("   ", base58.Encode(pubkey))
    }
    if seckey != nil && len(seckey) > 0 {
        xterm.Blue().Println("Encrpted Secret Key:")
        xterm.Yellow().Println("   ", base58.Encode(seckey))
    }
	xterm.Reset()

	return nil
}

//=============================================================================

func RunAsClient(raddr string) {

	ra, err := net.ResolveTCPAddr("tcp4", raddr)
	if err != nil {
		printErr(err, "Invalid remote address.")
		os.Exit(1)
	}

	g_chnl, err = xnet.Connect(ra, nil)
	if err != nil {
		printErr(err, "Can not connect to the keytool server.")
		os.Exit(1)
	}
	defer g_chnl.Close()

	ln := liner.NewLiner()
	ln.SetCtrlCAborts(true)
	defer ln.Close()

	handlers := make(map[string]func(int, []string) error)
	handlers["genkey"] = genkey
	handlers["getkey"] = getkey

	for {
		cmd, err := ln.Prompt("> ")
		if err != nil {
			break
		}
		ln.AppendHistory(cmd)

		if cmd == "quit" || cmd == "exit" {
			break
		}

		words := strings.Split(strings.TrimSpace(cmd), " ")
		handler, ok := handlers[words[0]]
		if !ok {
			xterm.Red().Println("Error: Unknown command.\n").Reset()
			continue
		}

		if err = handler(RunMode_Client, words); err != nil {
			xterm.Red().Printf("Error: %v\n\n", err).Reset()
			continue
		}

		fmt.Println()
	}
}
